import pymysql
import json
from typing import Dict, Any
from langchain_mcp_adapters.client import MultiServerMCPClient
from .db_config import get_db_config
from functools import lru_cache

# config 缓存
@lru_cache(maxsize=16)
def get_mcp_servers_config_by_domain(domain: str) -> Dict[str, Dict[str, Any]]:
    """
    根据 domain 查询数据库，返回该领域关联的所有 mcp_server 配置。
    :param domain: 领域ID（如 'law', 'medical'）
    :return: mcp_servers_config dict
    """
    db_config = get_db_config()
    conn = pymysql.connect(**db_config)
    try:
        with conn.cursor() as cursor:
            # 1. 查 domain_server 关联的 server_id
            cursor.execute("SELECT server_id FROM domain_server WHERE domain_id=%s", (domain,))
            server_ids = [str(row['server_id']) for row in cursor.fetchall()]
            if not server_ids:
                return {}
            # 2. 查 mcp_server 配置
            format_strings = ','.join(['%s'] * len(server_ids))
            cursor.execute(f"SELECT id, name, config FROM mcp_server WHERE id IN ({format_strings})", tuple(server_ids))
            servers = cursor.fetchall()
            mcp_servers_config = {}
            for server in servers:
                config_raw = server.get('config')
                if not config_raw:
                    continue
                try:
                    config_obj = json.loads(config_raw)
                    # 兼容 config 里有 mcpServers 多实例的情况
                    if 'mcpServers' in config_obj:
                        for k, v in config_obj['mcpServers'].items():
                            mcp_servers_config[k] = v
                    else:
                        mcp_servers_config[server['name']] = config_obj
                except Exception as e:
                    continue
            return mcp_servers_config
    finally:
        conn.close()

# tools 缓存（异步，需用全局字典）
_tools_cache = {}

# 新增：根据domain返回工具集
async def get_tools_by_domain(domain: str):
    """
    根据 domain 返回该领域的所有 MCP 工具集。
    :param domain: 领域ID（如 'law', 'medical'）
    :return: 工具列表
    """
    if domain in _tools_cache:
        return _tools_cache[domain]
    config = get_mcp_servers_config_by_domain(domain)
    mcp_client = MultiServerMCPClient(config)
    tools = await mcp_client.get_tools()
    _tools_cache[domain] = tools
    return tools 